_base_ = ['co_dino_5scale_r50_8xb2_1x_coco.py']

# pretrained = 'https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window12_384_22k.pth'  # noqa
# load_from = 'https://download.openmmlab.com/mmdetection/v3.0/codetr/co_dino_5scale_swin_large_16e_o365tococo-614254c9.pth'  # noqa
# load_from = 'work_dirs/co_dino_5scale_swin_l_16xb1_16e_o365tococo_package_aug/epoch_6.pth'

# model settings
# model = dict(
#     backbone=dict(
#         _delete_=True,
#         type='SwinTransformer',
#         pretrain_img_size=384,
#         embed_dims=192,
#         depths=[2, 2, 18, 2],
#         num_heads=[6, 12, 24, 48],
#         window_size=12,
#         mlp_ratio=4,
#         qkv_bias=True,
#         qk_scale=None,
#         drop_rate=0.,
#         attn_drop_rate=0.,
#         drop_path_rate=0.3,
#         patch_norm=True,
#         out_indices=(0, 1, 2, 3),
#         # Please only add indices that would be used
#         # in FPN, otherwise some parameter will not be used
#         with_cp=True,
#         convert_weights=True,
#         # init_cfg=dict(type='Pretrained', checkpoint=pretrained)
#         ),
#     neck=dict(in_channels=[192, 384, 768, 1536]),
#     query_head=dict(
#         dn_cfg=dict(box_noise_scale=0.4, group_cfg=dict(num_dn_queries=500)),
#         transformer=dict(encoder=dict(with_cp=6))))
model = dict(
    backbone=dict(
        _delete_=True,
        type='SwinTransformer',
        pretrain_img_size=384,
        # embed_dims=192,
        embed_dims=96,
        depths=[2, 2, 18, 2],
        # num_heads=[6, 12, 24, 48],
        num_heads=[3, 6, 12, 24],
        # window_size=12,
        window_size=16,
        mlp_ratio=4,
        qkv_bias=True,
        qk_scale=None,
        drop_rate=0.,
        attn_drop_rate=0.,
        drop_path_rate=0.3,
        patch_norm=True,
        out_indices=(0, 1, 2, 3),
        # Please only add indices that would be used
        # in FPN, otherwise some parameter will not be used
        with_cp=True,
        convert_weights=True,
        init_cfg=dict(type='Pretrained', checkpoint=pretrained)),
    # neck=dict(in_channels=[192, 384, 768, 1536]),
    neck=dict(in_channels=[96, 192, 384, 768]),
    query_head=dict(
        dn_cfg=dict(box_noise_scale=0.4, group_cfg=dict(num_dn_queries=500)),
        transformer=dict(encoder=dict(with_cp=6))))

image_size = (480, 640)

train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(type='PhotoMetricDistortion', hue_delta=36,),
    dict(type='RandomFlip', prob=0.5),
    dict(
        type='RandomChoice',
        transforms=[
            [
                dict(
                    type='Resize',
                    scale=image_size,#(720, 960)
                    keep_ratio=True)
            ],
            [
                dict(
                    type='RandomResize',
                    scale=image_size,
                    ratio_range=(1,2),
                    keep_ratio=True),
                dict(
                    type='MinIoURandomCrop',
                    ),
                dict(
                    type='Resize',
                    scale=image_size,
                    keep_ratio=True)
            ]
        ]),
    dict(type='PackDetInputs')
]

metainfo = {
    'classes': ('0', '1', '2', '3', '4', '5', '6', '7', '8', '9' ),
}

train_dataloader = dict(
    batch_size=4, num_workers=1, 
    dataset=dict(
        _delete_=True,
        type=_base_.dataset_type,
        data_root='/home/ma-user/work/mmdetection',#_base_.data_root,
        ann_file='raw_trainingset_coco/labels.json',
        # ann_file='ng_trainingset_coco/labels.json',
        data_prefix=dict(img='/'),
        filter_cfg=dict(filter_empty_gt=False, min_size=32),
        pipeline=train_pipeline,
        backend_args=_base_.backend_args,
        metainfo=metainfo
    )
)

test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='Resize', scale=image_size, keep_ratio=True),
    dict(type='LoadAnnotations', with_bbox=True),
    dict(
        type='PackDetInputs',
        meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape',
                   'scale_factor'))
]

val_dataloader = dict(
    batch_size=1, num_workers=1, 
    dataset=dict(
        _delete_=True,
        type=_base_.dataset_type,
        data_root='/home/ma-user/work/mmdetection',#_base_.data_root,
        ann_file='raw_validset_coco/labels.json',
        data_prefix=dict(img='/'),
        filter_cfg=dict(filter_empty_gt=False, min_size=32),
        pipeline=test_pipeline,
        backend_args=_base_.backend_args,
        metainfo=metainfo
    )
)

test_dataloader = dict(
    batch_size=4, num_workers=1, 
    dataset=dict(
        _delete_=True,
        type=_base_.dataset_type,
        data_root='/home/ma-user/work/mmdetection',#_base_.data_root,
        ann_file='raw_testset_coco/labels.json',
        data_prefix=dict(img='/'),
        filter_cfg=dict(filter_empty_gt=False, min_size=32),
        pipeline=test_pipeline,
        backend_args=_base_.backend_args,
        metainfo=metainfo
    )
)

val_evaluator = dict(ann_file='/home/ma-user/work/mmdetection/raw_validset_coco/labels.json')
test_evaluator = dict(ann_file='/home/ma-user/work/mmdetection/raw_testset_coco/labels.json')


optim_wrapper = dict(optimizer=dict(lr=1e-4))

max_epochs = 16
train_cfg = dict(max_epochs=max_epochs)

param_scheduler = [
    dict(
        type='MultiStepLR',
        begin=0,
        end=max_epochs,
        by_epoch=True,
        milestones=[8],
        gamma=0.1)
]

default_hooks = dict(
    checkpoint=dict(by_epoch=True, interval=1, max_keep_ckpts=16))

batch_augments = [
    dict(type='BatchFixedSizePad', size=image_size, pad_mask=True)
]